We consider the problem of estimating the input and hidden variables of a stochastic multi-layer neural network from an observation of the output. The hidden variables in each layer are represented as matrices with statistical interactions along both rows as well as columns. This problem applies to matrix imputation, signal recovery via deep generative prior models, multi-task and mixed regression, and learning certain classes of two-layer neural networks. We extend a recently-developed algorithm -- Multi-Layer Vector Approximate Message Passing (ML-VAMP), for this matrix-valued inference problem. It is shown that the performance of the proposed Multi-Layer Matrix VAMP (ML-Mat-VAMP) algorithm can be exactly predicted in a certain random large-system limit, where the dimensions $N\times d$ of the unknown quantities grow as $N\rightarrow\infty$ with $d$ fixed. In the two-layer neural-network learning problem, this scaling corresponds to the case where the number of input features, as well as training samples, grow to infinity but the number of hidden nodes stays fixed. The analysis enables a precise prediction of the parameter and test error of the learning.

Summary and Contributions: The authors extend the Multi Layer Vector Approximate Message Passing (ML-VAMP) framework, used for inference of signals in multi-layer Generalised Linear Models, to the matrix-valued signal case, and call the algorithm ML-Mat-VAMP. In addition they provide an asymptotic analysis through the demonstration of so-called state evolution equations, that allow to describe the dynamics of the ML-Mat-VAMP through "macroscopic variables" such as mean-square error etc. They also provide numerical experiments to illustrate the validity of the state evolution in order to test error in simple learning model of a shallow neural net. UPDATE POST AUTHOR RESPONSE: My concerns have been mostly answered. About the main concern: the authors mention they will update the numerical part so that their example will not be covered anymore by the existing theory found in ref [1] which is satisfying.

We consider the problem of estimating the input and hidden variables of a stochastic multi-layer neural network from an observation of the output. The hidden variables in each layer are represented as matrices with statistical interactions along both rows as well as columns. This problem applies to matrix imputation, signal recovery via deep generative prior models, multi-task and mixed regression, and learning certain classes of two-layer neural networks. We extend a recently-developed algorithm -- Multi-Layer Vector Approximate Message Passing (ML-VAMP), for this matrix-valued inference problem. It is shown that the performance of the proposed Multi-Layer Matrix VAMP (ML-Mat-VAMP) algorithm can be exactly predicted in a certain random large-system limit, where the dimensions N\times d of the unknown quantities grow as N\rightarrow\infty with d fixed.